/*=========================================================================
| (c) 2011  Total Phase, Inc.
|--------------------------------------------------------------------------
| Project : Cheetah SST Programming
| File    : sst_spi.c
|--------------------------------------------------------------------------
| Read from, erase, write to, and verify data contents of an SST Flash Chip
|--------------------------------------------------------------------------
| Redistribution and use of this file in source and binary forms, with
| or without modification, are permitted.
|
| THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
| "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
| LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
| FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
| COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
| INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
| BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
| LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
| CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
| LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
| ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
| POSSIBILITY OF SUCH DAMAGE.
 ========================================================================*/

//=========================================================================
// INCLUDES
//=========================================================================
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <libgen.h>

#include "cheetah.h"

#ifdef _WIN32
#include <time.h>
#else
#include <sys/time.h>
#endif


//=========================================================================
// SST CONSTANTS
//=========================================================================
#define MEMORY_SIZE (512*1024) // 512 KB
#define SECTOR_SIZE (64*1024)  // 64 KB
#define NUM_SECTORS (MEMORY_SIZE / SECTOR_SIZE)

#define READ_BLOCK_SIZE (8*1024)

#define PAGE_SIZE                   10240
#define PAGE_write_BATCH_SIZE       1
#define WORD_PROGRAM_TIME_NS        10000

#define INSTR_FAST_READ     0x0b
#define INSTR_write_ENABLE  0x06
#define INSTR_write_DISABLE 0x04
#define INSTR_PAGE_PROGRAM  0x02
#define INSTR_AAI_PROGRAM   0xAD
#define INSTR_SECTOR_ERASE  0xd8
#define INSTR_BULK_ERASE    0xc7
#define INSTR_READ_STATUS   0x05
#define INSTR_write_STATUS  0x01
#define INSTR_IDCODE        0x9F


//=========================================================================
// PROGRAM CONSTANTS
//=========================================================================
#define COMMAND_READ    0
#define COMMAND_write   1
#define COMMAND_ERASE   2
#define COMMAND_EWRITE  4
#define COMMAND_VERIFY  5
#define COMMAND_IDCODE  6

#define SS_OFF    0x00
#define SS_ASSERT 0x01


//=========================================================================
// FUNCTIONS
//=========================================================================
static s64 _timeMicroseconds () {
#ifdef _WIN32
    return ((s64)clock()) * 1000000 / CLOCKS_PER_SEC;
#else
    struct timeval tv;
    gettimeofday(&tv, 0);
    return ((s64)tv.tv_sec * 1000000L) + (s64)(tv.tv_usec);
#endif
}

static void _printUsage (char *progname) {
    printf("%s - utility for programming SST SPI Flash\n",
           progname);
    printf("usage: %s PORT BITRATE read   LEN_KB\n", progname);
    printf("usage: %s PORT BITRATE write  FILENAME\n", progname);
    printf("usage: %s PORT BITRATE erase  SECTOR NUM\n", progname);
    printf("usage: %s PORT BITRATE ewrite FILENAME\n", progname);
    printf("usage: %s PORT BITRATE verify FILENAME\n", progname);
    printf("usage: %s PORT BITRATE idcode\n", progname);
    printf("\n");
}

static void _printProgress (int percent, double elapsedTime, int pos)
{
    static char progressbar[32];
    int tenths = percent/5 + 1;
    sprintf(progressbar, "[%*c%*c", tenths, '#', 22-tenths, ']');
    printf("\r%s %3d%% %7d KB %.2lf seconds",
           progressbar, percent,
           pos / 1024,
           elapsedTime);
    fflush(stdout);
}

static int _connect (int port, int bitrate, int mode, Cheetah *handle) {
    // Open the device.
    *handle = ch_open(port);
    if (*handle <= 0) {
        printf("Unable to open Cheetah device on port %d\n", port);
        printf("Error code = %d\n", *handle);
        return 1;
    }

    printf("Opened Cheetah device on port %d\n", port);

    printf("Host interface is %s\n",
           (ch_host_ifce_speed(*handle)) ? "high speed" : "full speed");

    // Ensure that the SPI subsystem is configured.
    ch_spi_configure(*handle, (mode >> 1), mode & 1,
                     CH_SPI_BITORDER_MSB, 0x00);
    printf("SPI configuration set to mode %d, MSB shift, "
           "SS[2:0] polarity = 0x%01x\n",
           mode, 0x00);
    fflush(stdout);

    // Turn target power on
    ch_target_power(*handle, CH_TARGET_POWER_ON);
    ch_sleep_ms(100);

    // Set the bitrate.
    bitrate = ch_spi_bitrate(*handle, bitrate);
    if (bitrate < 0) {
        fprintf(stderr, "error: could not set bitrate.\n");
        fprintf(stderr, "  Error code = %d\n", bitrate);
        return 1;
    }
    printf("Bitrate set to %d kHz\n", bitrate);

    // Enable the SPI outputs
    ch_spi_queue_oe(*handle, 1);

    fflush(stdout);
    return 0;
}


int _do_simple_inst (Cheetah handle,   u08 instruction,
                     int     num_read, u08 *data_in)
{
    ch_spi_queue_clear(handle);
    ch_spi_queue_ss(handle, SS_ASSERT);
    ch_spi_queue_byte(handle, 1, instruction);
    ch_spi_queue_byte(handle, num_read, 0);
    ch_spi_queue_ss(handle, SS_OFF);

    if (num_read) {
        int len = num_read+1;
        int ret = ch_spi_batch_shift(handle, len, data_in);
        if (ret != len) {
            fprintf(stderr,
                    "error: expected %d bytes but only received %d bytes\n",
                    len, ret);
            return -1;
        }
        return len;
    }
    else {
        ch_spi_batch_shift(handle, 0, 0);
        return 0;
    }
}


void _clear_block_protect (Cheetah handle) {
    // Queue the write enable instruction for the flash
    // and clear the block protection bits
    ch_spi_queue_ss(handle, SS_ASSERT);
    ch_spi_queue_byte(handle, 1, INSTR_WRITE_ENABLE);
    ch_spi_queue_ss(handle, SS_OFF);
    ch_spi_queue_delay_cycles(handle, 8);

    ch_spi_queue_ss(handle, SS_ASSERT);
    ch_spi_queue_byte(handle, 1, INSTR_WRITE_STATUS);
    ch_spi_queue_byte(handle, 1, 0x00);
    ch_spi_queue_ss(handle, SS_OFF);
    ch_spi_batch_shift(handle, 0, 0);
}

int _read_status_reg (Cheetah handle) {
    u08 data[1+1];

    int ret = _do_simple_inst(handle, INSTR_READ_STATUS, 1, data);
    if (ret < 0)
        return ret;
    return data[1];
}

int _read_mem (Cheetah handle, u08 *data_in, int length, int blocksize)
{
    int count;
    s64 start, currTime;
    int ret = 0;
    int pos = 0;
    int total_length = length;

    if (length == 0 || blocksize == 0)
        return 0;

    ch_spi_queue_clear(handle);
    ch_spi_queue_ss(handle, SS_ASSERT);

    // Queue fast read command code.
    ch_spi_queue_byte(handle, 1, INSTR_FAST_READ);

    // Queue 3 bytes of address.
    ch_spi_queue_byte(handle, 1, 0);
    ch_spi_queue_byte(handle, 1, 0);
    ch_spi_queue_byte(handle, 1, 0);

    // Queue dummy byte.
    ch_spi_queue_byte(handle, 1, 0);

    // Shift the queued fast read command.
    count = ch_spi_batch_shift(handle, 5, data_in);
    if (count != 5) {
        fprintf(stderr, "error: expected 5 bytes from initial shift\n");
        ret = 2;
        goto _read_mem_quit;
    }

    start = _timeMicroseconds();
    _printProgress(0, 0, 0);

    // Read one block at a time.
    while (length > 0) {
        int percent;
        int iter_len = (length > blocksize) ? blocksize : length;

        // Read the next block of data.
        ch_spi_queue_clear(handle);
        ch_spi_queue_byte(handle, iter_len, 0x00);
        count = ch_spi_batch_shift(handle, iter_len, data_in);
        if (count != iter_len) {
            fprintf(stderr,
                    "error: expected %d bytes from block shift (got %d)\n",
                    iter_len, count);
            ret = 2;
            break;
        }

        // Print out the progress.
        currTime = _timeMicroseconds();

        data_in += iter_len;
        length  -= iter_len;
        pos     += iter_len;

        if (total_length < 100)
            percent = (pos * 100) / total_length;
        else
            // Prevent overflow by rearranging equation
            percent = (pos) / (total_length / 100);

        _printProgress(percent,
                       ((double)(currTime - start)) / 1000000,
                       pos);
    }

_read_mem_quit:
    return ret;
}

static int read_helper (Cheetah handle, int length) {
    int  ret = 0;
    u08 *data_in;
    int  i;

    // Convert address and length from KB to bytes
    length *= 1024;

    if (length > MEMORY_SIZE) {
        printf("warning: read length truncated to memory size\n");
        length = MEMORY_SIZE;
    }

    // Create the buffer to receive the flash values into.
    data_in = (u08 *)malloc(length);
    memset(data_in, 0, length);

    ret = _read_mem(handle, data_in, length, READ_BLOCK_SIZE);
    if (ret)
        goto _read_quit;

    printf("\n");

    // Dump the data to the screen
    printf("\nData read from device:");
    for (i=0; i < length; ++i) {
        u08 c = data_in[i];

        if ((i&0x0f) == 0)      printf("\n%04x:  ", i);
        printf("%02x ", c & 0xff);
        if (((i+1)&0x07) == 0)  printf(" ");
    }
    printf("\n");
    fflush(stdout);

_read_quit:
    // Free the packet
    free(data_in);

    return ret;
}

static int _verify (Cheetah handle, u08 *prog_data, int prog_len)
{
    int  ret = 0;
    int  i;
    u08 *data_in;

    // Variables for the integrity check.
    int integrity_errors = 0;

    // Create the buffer to receive the flash values into.
    data_in = (u08 *)malloc(prog_len);

    ret =_read_mem(handle, data_in, prog_len, READ_BLOCK_SIZE);
    if (ret)
        goto _verify_quit;

    for (i=0; i < prog_len; ++i) {
        u08 c = data_in[i];
        if (c != prog_data[i]) {
            printf("in: %02x, out: %02x\n", data_in[i], prog_data[i]);
            ++integrity_errors;
        }
    }

    fflush(stdout);
    printf("\nThere were %d data errors.\n", integrity_errors);

    // Release the buffer.
_verify_quit:
    free(data_in);

    if (integrity_errors)
        ret = 2;

    return ret;
}

static int write_helper (Cheetah handle, u08 *prog_data, int prog_len) {
    int ret = 0;

    _clear_block_protect(handle);

    printf("Writing to flash:\n");
    s64 start = _timeMicroseconds();
    _printProgress(0, 0, 0);

    _do_simple_inst(handle, INSTR_WRITE_ENABLE, 0, 0);

    int addr = 0;
    int  length   = prog_len;
    u08 *data_out = prog_data;

    ch_spi_queue_clear(handle);
    ch_spi_queue_ss(handle, SS_ASSERT);
    ch_spi_queue_byte(handle, 1, INSTR_AAI_PROGRAM);
    ch_spi_queue_byte(handle, 1, (addr >> 16) & 0xff);
    ch_spi_queue_byte(handle, 1, (addr >>  8) & 0xff);
    ch_spi_queue_byte(handle, 1, (addr >>  0) & 0xff);

    while (length > 0) {
        // Check if we've reached the end.
        int page_size = (length > PAGE_SIZE) ? PAGE_SIZE : length;
        int j;
        for (j = 0; j < page_size; j += 2, addr += 2) {
            if (addr > 0) {
                ch_spi_queue_delay_ns(handle, WORD_PROGRAM_TIME_NS);
                ch_spi_queue_ss(handle, SS_ASSERT);
                ch_spi_queue_byte(handle, 1, INSTR_AAI_PROGRAM);
            }
            ch_spi_queue_byte(handle, 1, data_out[j]);
            ch_spi_queue_byte(handle, 1, data_out[j+1]);
            ch_spi_queue_ss(handle, SS_OFF);
        }

        data_out += page_size;
        length   -= page_size;

        // Shift out the write command.  (Don't need the data back.)
        int batch = ch_spi_batch_length(handle);
        int count = ch_spi_batch_shift(handle, 0, 0);
        if (count != batch) {
            printf("Expected %d bytes but only received %d bytes\n",
                   batch, count);
            ret = 2;
            break;
        }
        ch_spi_queue_clear(handle);

        _read_status_reg(handle);

        // Print out the progress.
        s64 currTime = _timeMicroseconds();

        int percent;
        if (prog_len < 100)
            percent = (addr * 100) / prog_len;
        else
            // Prevent overflow by rearranging equation
            percent = addr / (prog_len / 100);

        _printProgress(percent,
                       ((double)(currTime - start)) / 1000000,
                       addr);
    }

    printf("\n");

    _do_simple_inst(handle, INSTR_WRITE_DISABLE, 0, 0);
    int status_reg = _read_status_reg(handle);
    printf("status reg = 0x%02x\n", status_reg);

    if (!ret) {
        printf("Verifying:\n");
        _verify(handle, prog_data, prog_len);
    }

    return ret;
}

static int _erase (Cheetah handle, int sector, int num) {
    int   ret = 0;
    int   addr;
    int   eraseAll = 0;
    int   batch, count;
    char *str;
    s64   start;
    s64   end;

    if (sector + num > NUM_SECTORS) {
        printf("warning: erase sectors truncted to size of device\n");
        num = NUM_SECTORS - sector;
    }

    if (sector < 0 || sector >= NUM_SECTORS) {
        printf("erase start sector is invalid\n");
        return 1;
    }

    if (sector == 0 && num == NUM_SECTORS)
        eraseAll = 1;

    if (eraseAll)  str = "Bulk";
    else           str = "Block";

    // Reset the state of the bus.
    _clear_block_protect(handle);
    printf("status reg = 0x%02x\n", _read_status_reg(handle));

    while (num) {
        addr = sector << 16;
        if (!eraseAll) {
            printf("Erasing sector %02d (bytes 0x%06x to 0x%06x)...\n",
                   sector, addr, addr | 0xffff);
            fflush(stdout);
        }
        else {
            printf("Erasing entire device...\n");
            fflush(stdout);
        }

        // Start the erase sequence.
        ch_spi_queue_clear(handle);

        // Queue the write enable instruction for the flash.
        ch_spi_queue_ss(handle, SS_ASSERT);
        ch_spi_queue_byte(handle, 1, INSTR_WRITE_ENABLE);
        ch_spi_queue_ss(handle, SS_OFF);

        ch_spi_queue_ss(handle, SS_ASSERT);

        if (!eraseAll) {
            // Queue the sector erase command.
            ch_spi_queue_byte(handle, 1, INSTR_SECTOR_ERASE);
            ch_spi_queue_byte(handle, 1, (u08)((addr >> 16) & 0xff));
            ch_spi_queue_byte(handle, 1, (u08)((addr >>  8) & 0xff));
            ch_spi_queue_byte(handle, 1, (u08)((addr >>  0) & 0xff));
        }
        else {
            // Queue the bulk erase command.
            ch_spi_queue_byte(handle, 1, INSTR_BULK_ERASE);
        }

        ch_spi_queue_ss(handle, SS_OFF);
        // Shift the queued commands.  (Don't need the data back.)
        batch = ch_spi_batch_length(handle);
        count = ch_spi_batch_shift(handle, 0, 0);
        if (count != batch) {
            fprintf(stderr,
                    "error: expected %d bytes but only received %d bytes\n",
                   batch, count);
            ret = 2;
            break;
        }

        start = _timeMicroseconds();
        while (1) {
            u08 status_reg = _read_status_reg(handle);
            printf("status_reg = 0x%02x\n", status_reg);
            if (!(status_reg & 0x01))
                break;

            ch_sleep_ms(10);
        }
        end = _timeMicroseconds();
        printf("%s erase took %.03f seconds\n",
               str, (double)(end-start)/1000000);
        fflush(stdout);

        if (!eraseAll) {
            ++sector;
            --num;
        }
        else {
            sector += 64;
            num     = 0;
        }
    }

    return ret;
}


static int _idcode (Cheetah handle) {
    int ret = 0;

    u08 data[1+3];
    // Queue the write enable instruction for the flash.
    ret = _do_simple_inst(handle, INSTR_IDCODE, 3, data);
    if (ret < 0) {
        ret = 2;
        goto done;
    }
    u32 idcode = (data[1] << 16) | (data[2] << 8) | data[3];
    printf("idcode = 0x%06x\n", idcode);

    // Queue the write enable instruction for the flash.
    int status_reg = _read_status_reg(handle);
    if (status_reg < 0) {
        ret = 2;
        goto done;
    }
    printf("status reg = 0x%02x\n", status_reg);

done:
    return ret;
}


//=========================================================================
// MAIN PROGRAM
//=========================================================================
int main (int argc, char *argv[]) {
    int         ret = 0;

    char        progname[256];

    Cheetah     handle;
    int         port    = 0;
    int         bitrate = 0;
    int         mode    = 0;

    const char *command;
    int         commandID = -1;

    u08        *prog_data = 0;
    int         length = 0;
    int         sector = 0;

    strcpy(progname, basename(argv[0]));

    // Print the syntax for usage.
    if (argc < 4) {
        _printUsage(progname);
        return 1;
    }

    // Parse the arguments.
    port      = atoi(argv[1]);
    bitrate   = atoi(argv[2]);
    command   = argv[3];

    // Determine which command transaction type was requested.
    if (strcmp(command, "read") == 0) {
        commandID = COMMAND_READ;
        if (argc < 5) {
            _printUsage(progname);
            return 1;
        }
    }
    else if (strcmp(command, "write") == 0) {
        if (argc < 5) {
            _printUsage(progname);
            return 1;
        }

        commandID = COMMAND_WRITE;
    }
    else if (strcmp(command, "erase") == 0) {
        commandID = COMMAND_ERASE;
        if (argc <= 4) {
            sector = 0;
            length = NUM_SECTORS;
        }
        else if (argc < 6)
        {
            _printUsage(progname);
            return 1;
        }
        else {
            sector = atoi(argv[4]);
            length = atoi(argv[5]);
        }
    }
    else if (strcmp(command, "ewrite") == 0) {
        if (argc < 5) {
            _printUsage(progname);
            return 1;
        }

        commandID = COMMAND_EWRITE;
    }
    else if (strcmp(command, "verify") == 0) {
        if (argc < 5) {
            _printUsage(progname);
            return 1;
        }

        commandID = COMMAND_VERIFY;
    }
    else if (strcmp(command, "idcode") == 0) {
        commandID = COMMAND_IDCODE;
    }
    else {
        fprintf(stderr, "error: unknown option: %s\n", command);
        fprintf(stderr,
                "Valid options are: read, write, erase, ewrite and verify\n\n");
        _printUsage(progname);
        return 1;
    }

    if (commandID == COMMAND_WRITE  ||
        commandID == COMMAND_EWRITE ||
        commandID == COMMAND_VERIFY)
    {
        char *filename = argv[4];
        FILE *handle = fopen(filename, "rb");
        struct stat buf;
        int ret = 0;

        if (!handle) {
            fprintf(stderr,
                    "error: could not open programming file '%s' for reading\n",
                    filename);
            return 1;
        }

        fstat(fileno(handle), &buf);
        length = buf.st_size;
        if (length > MEMORY_SIZE) {
            printf("warning: programming file truncated to memory size\n");
            length = MEMORY_SIZE;
        }

        prog_data = (u08 *)malloc(length);
        ret = fread(prog_data, length, 1, handle);
        if (ret != 1) {
            fprintf(stderr,
                    "error: could not read programming data from file\n");
            return 1;
        }

        fclose(handle);
    }

    // Connect to and configure the Cheetah.
    mode = 0;
    if (_connect(port, bitrate, mode, &handle)) {
        return 1;
    }

    // Execute the appropriate command.
    switch (commandID) {
      case COMMAND_READ:
        ret = read_helper(handle, atoi(argv[4]));
        break;

      case COMMAND_WRITE:
        ret = write_helper(handle, prog_data, length);
        break;

      case COMMAND_ERASE:
        ret = _erase(handle, sector, length);
        break;

      case COMMAND_EWRITE:
        ret = _erase(handle, 0, NUM_SECTORS);
        if (!ret)
            write_helper(handle, prog_data, length);
        break;

      case COMMAND_VERIFY:
        ret = _verify(handle, prog_data, length);
        break;

      case COMMAND_IDCODE:
        ret = _idcode(handle);
        break;
    }

    // Close the device.
    ch_target_power(handle, CH_TARGET_POWER_OFF);
    ch_spi_queue_oe(handle, 0);
    ch_close(handle);

    if (prog_data)
        free(prog_data);

    return 0;
}
